#include "plato_variable.hh"
#include "../impl/plato_stream_impl.hh"

plato::StructVariable::StructVariable(Domain *domain, VarID parent,
                                      PlatoVariableSyncType sync_type)
    : VariableImpl(domain, PlatoType::STRUCT, parent, sync_type) {}

plato::StructVariable::StructVariable(const std::string &name, Domain *domain,
                                      VarID parent,
                                      PlatoVariableSyncType sync_type)
    : VariableImpl(name, domain, PlatoType::STRUCT, parent, sync_type) {}

plato::StructVariable::~StructVariable() {}

auto plato::VariableImpl::set_id(VarID id) -> void { var_id_ = id; }

plato::VariableImpl::VariableImpl(Domain *domain, PlatoType type, VarID parent,
                                  PlatoVariableSyncType sync_type) {
  domain_ = domain;
  type_ = type;
  parent_ = parent;
  sync_type_ = sync_type;
}

plato::VariableImpl::VariableImpl(const std::string &name, Domain *domain,
                                  PlatoType type, VarID parent,
                                  PlatoVariableSyncType sync_type) {
  name_ = name;
  domain_ = domain;
  type_ = type;
  parent_ = parent;
  sync_type_ = sync_type;
}

plato::VariableImpl::~VariableImpl() { domain_->remove(var_id_); }

auto plato::VariableImpl::type() const -> PlatoType { return type_; }

auto plato::VariableImpl::name() -> const std::string & { return name_; }

auto plato::VariableImpl::parent() const -> VariablePtr {
  return domain_->get(parent_);
}

auto plato::VariableImpl::domain() const -> Domain * { return domain_; }

auto plato::VariableImpl::sync_type() const -> PlatoVariableSyncType {
  return sync_type_;
}

auto plato::VariableImpl::id() const -> VarID { return var_id_; }

namespace plato {

class DomainImpl : public Domain {
  VarID var_id_{1};
  using VarMap = std::unordered_map<VarID, VariablePtr>;
  using VarNameMap = std::unordered_map<std::string, VariablePtr>;
  VarMap var_map_;
  VarNameMap var_name_map_;
  PlatoStreamImpl stream_;
  DomainID id_{INVALID_DOMAIN_ID};
  MemBlock *mb_{nullptr};
  bool mb_delete_{false};

public:
  DomainImpl( DomainID id = INVALID_DOMAIN_ID, MemBlock *mb = nullptr);
  virtual ~DomainImpl();
  virtual auto identify(VariablePtr ptr) -> bool override;
  virtual auto remove(VarID id) -> void override;
  virtual auto get(VarID id) -> VariablePtr override;
  virtual auto get(const std::string &name) -> VariablePtr override;
  virtual auto get_stream() -> PlatoStream & override;
  virtual auto do_sync() -> void override;
  virtual auto serialize_all() -> void override;
  virtual auto id() -> DomainID override;
  virtual auto mem_block() -> MemBlock * override;
  virtual auto reset_named_var() -> void override {
    for (auto &[_, v] : var_name_map_) {
      v->copy_default();
    }
  }
};

StringVariable::StringVariable(Domain *domain, VarID parent,
                               PlatoVariableSyncType sync_type)
    : VariableImpl(domain, PlatoType::STRING, parent, sync_type) {}

StringVariable::StringVariable(const std::string &name, Domain *domain,
                               VarID parent, PlatoVariableSyncType sync_type)
    : VariableImpl(name, domain, PlatoType::STRING, parent, sync_type) {}

StringVariable::StringVariable(Domain *domain, VarID parent,
                               PlatoVariableSyncType sync_type,
                               const std::string &t)
    : VariableImpl(domain, PlatoType::STRING, parent, sync_type) {
  s_ = t;
}

StringVariable::StringVariable(const std::string &name, Domain *domain,
                               VarID parent, PlatoVariableSyncType sync_type,
                               const std::string &t)
    : VariableImpl(name, domain, PlatoType::STRING, parent, sync_type) {
  s_ = t;
}

StringVariable::~StringVariable() {}

auto StringVariable::New(Domain *domain, VarID parent,
                         PlatoVariableSyncType sync_type)
    -> std::shared_ptr<StringVariable> {
  return std::make_shared<StringVariable>(domain, parent, sync_type);
}

auto StringVariable::serialize(PlatoStream &stream) -> bool {
  if (sync_type() == PlatoVariableSyncType::SENDER) {
    stream << id();
    stream << (std::uint32_t)s_.size();
    stream.write(s_.data(), s_.size());
    return true;
  }
  return false;
}

auto StringVariable::deserialize(PlatoStream &stream) -> bool {
  if (sync_type() == PlatoVariableSyncType::RECEIVER) {
    VarID var_id;
    stream >> var_id;
    set_id(var_id);
    std::uint32_t size = 0;
    stream >> size;
    stream.read(get_buffer(size), size);
    s_.assign(get_buffer(), size);
    return true;
  }
  return false;
}

StringVariable::operator std::string() { return s_; }

BoolVariable::BoolVariable(Domain *domain, VarID parent,
                           PlatoVariableSyncType sync_type)
    : VariableImpl(domain, TypeValue<bool>::value, parent, sync_type) {}

BoolVariable::BoolVariable(const std::string &name, Domain *domain,
                           VarID parent, PlatoVariableSyncType sync_type)
    : VariableImpl(name, domain, TypeValue<bool>::value, parent, sync_type) {}

BoolVariable::BoolVariable(Domain *domain, VarID parent,
                           PlatoVariableSyncType sync_type, bool t)
    : VariableImpl(domain, TypeValue<bool>::value, parent, sync_type) {
  value_ = t;
}
BoolVariable::BoolVariable(const std::string &name, Domain *domain,
                           VarID parent, PlatoVariableSyncType sync_type,
                           bool t)
    : VariableImpl(name, domain, TypeValue<bool>::value, parent, sync_type) {
  value_ = t;
}

BoolVariable::~BoolVariable() {}

auto BoolVariable::New(Domain *domain, VarID parent,
                       PlatoVariableSyncType sync_type) {
  return std::make_shared<BoolVariable>(domain, parent, sync_type);
}

auto BoolVariable::serialize(PlatoStream &stream) -> bool {
  if (sync_type() == PlatoVariableSyncType::SENDER) {
    stream << id();
    stream << value_;
    return true;
  }
  return false;
}

auto BoolVariable::deserialize(PlatoStream &stream) -> bool {
  if (sync_type() == PlatoVariableSyncType::RECEIVER) {
    VarID var_id;
    stream >> var_id;
    set_id(var_id);
    stream >> value_;
    return true;
  }
  return false;
}

} // namespace plato

plato::DomainImpl::DomainImpl(DomainID id, MemBlock *mb) {
  id_ = id;
  if (mb) {
    mb_ = mb;
  } else {
    mb_delete_ = true;
    mb_ = new_mem_block_raw();
  }
}

plato::DomainImpl::~DomainImpl() {
  VarMap temp;
  temp.swap(var_map_);
  if (mb_delete_) {
    delete mb_;
  }
}

auto plato::DomainImpl::identify(VariablePtr ptr) -> bool {
  auto ptr_impl = std::dynamic_pointer_cast<VariableImpl>(ptr);
  if (!ptr_impl) {
    return false;
  }
  if (ptr_impl->id() == INVALID_VAR_ID) {
    ptr_impl->set_id(var_id_);
    var_map_.emplace(var_id_, ptr);
    var_id_ += 1;
  } else {
    if (ptr->id() > var_id_) {
      var_id_ = ptr->id() + 1;
    }
    var_map_.emplace(ptr->id(), ptr);
  }
  return true;
}

auto plato::DomainImpl::remove(VarID id) -> void { var_map_.erase(id); }

auto plato::DomainImpl::get(VarID id) -> VariablePtr {
  auto it = var_map_.find(id);
  if (it == var_map_.end()) {
    return nullptr;
  }
  return it->second;
}

auto plato::DomainImpl::get(const std::string &name) -> VariablePtr {
  auto it = var_name_map_.find(name);
  if (it == var_name_map_.end()) {
    return nullptr;
  }
  return it->second;
}

auto plato::DomainImpl::get_stream() -> PlatoStream & { return stream_; }

auto plato::DomainImpl::do_sync() -> void {
  while (stream_.available() > sizeof(VarID)) {
    VarID var_id = INVALID_VAR_ID;
    stream_.copy((char *)&var_id, sizeof(VarID));
    auto it = var_map_.find(var_id);
    if (it != var_map_.end()) {
      if (!it->second->deserialize(stream_)) {
        return;
      }
    }
  }
}

auto plato::DomainImpl::serialize_all() -> void {
  for (auto &[_, v] : var_map_) {
    v->serialize(stream_);
  }
}

auto plato::DomainImpl::id() -> DomainID { return id_; }

auto plato::DomainImpl::mem_block() -> MemBlock * { return mb_; }

static std::size_t buffer_size_ = 1024;
static std::shared_ptr<char[]> buffer_(new char[buffer_size_]);

char *plato::get_buffer(std::size_t size) {
  if (buffer_size_ < size) {
    buffer_ = std::shared_ptr<char[]>(new char[size]);
    buffer_size_ = size;
  }
  return buffer_.get();
}

char *plato::get_buffer() { return buffer_.get(); }

auto plato::new_domain(DomainID id, MemBlock* mb) -> DomainPtr {
  return std::make_shared<DomainImpl>(id, mb);
}
